#!/bin/bash

# Flink SQL HBase to HDFS Parquet Job Submission Script
# =====================================================

# 脚本配置
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
JOB_NAME="hbase2hdfs-parquet"
JAR_FILE="${PROJECT_ROOT}/target/flink-hbase-1.0-SNAPSHOT.jar"
MAIN_CLASS="com.flink.hbase.sql2hdfs.FlinkSqlHBaseToHdfsJob"

# 默认配置
DEFAULT_FLINK_HOME="/opt/flink"
DEFAULT_CONFIG_FILE="${PROJECT_ROOT}/src/main/resources/hbase2hdfs-config.properties"

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_debug() {
    echo -e "${BLUE}[DEBUG]${NC} $1"
}

# 打印使用说明
print_usage() {
    cat << EOF
Usage: $0 [OPTIONS]

Flink SQL HBase to HDFS Parquet Job Submission Script

OPTIONS:
    -h, --help                      显示帮助信息
    -c, --config FILE              配置文件路径 (默认: ${DEFAULT_CONFIG_FILE})
    -f, --flink-home PATH          Flink 安装路径 (默认: ${DEFAULT_FLINK_HOME})
    -j, --jar FILE                 JAR 文件路径 (默认: ${JAR_FILE})
    
    # HBase 配置
    --zk-quorum QUORUM             Zookeeper 集群地址 (如: zk1:2181,zk2:2181,zk3:2181)
    --table-name TABLE             HBase 表名
    --column-family FAMILY         HBase 列族名 (默认: info)
    
    # Kerberos 认证
    --enable-kerberos              启用 Kerberos 认证
    --principal PRINCIPAL          Kerberos 主体
    --keytab-path PATH             Keytab 文件路径
    --krb5-conf-path PATH          krb5.conf 文件路径
    
    # HDFS 配置
    --hdfs-output PATH             HDFS 输出路径
    --checkpoint-path PATH         Checkpoint 存储路径
    
    # 并行度配置
    --parallelism NUM              全局并行度 (默认: 4)
    --source-parallelism NUM       数据源并行度 (默认: 2)
    --sink-parallelism NUM         数据写入并行度 (默认: 2)
    
    # 其他选项
    --dry-run                      仅显示命令，不执行
    --verbose                      详细输出
    --validate                     验证环境和配置

EXAMPLES:
    # 基本使用
    $0 --zk-quorum zk1:2181,zk2:2181,zk3:2181 \\
       --table-name user_table \\
       --hdfs-output hdfs://namenode:9000/user/flink/output
    
    # 使用 Kerberos 认证
    $0 --zk-quorum zk1:2181 \\
       --table-name user_table \\
       --hdfs-output hdfs://namenode:9000/output \\
       --enable-kerberos \\
       --principal flink@EXAMPLE.COM \\
       --keytab-path /etc/security/keytabs/flink.keytab
    
    # 使用配置文件
    $0 -c hbase2hdfs-config.properties
    
    # 验证环境
    $0 --validate

EOF
}

# 检查环境
check_environment() {
    log_info "检查环境配置..."
    
    # 检查 Flink 安装
    if [[ ! -d "$FLINK_HOME" ]]; then
        log_error "Flink 安装目录不存在: $FLINK_HOME"
        log_info "请设置 FLINK_HOME 环境变量或使用 --flink-home 参数"
        exit 1
    fi
    
    # 检查 flink 命令
    if [[ ! -x "$FLINK_HOME/bin/flink" ]]; then
        log_error "Flink 命令不存在: $FLINK_HOME/bin/flink"
        exit 1
    fi
    
    # 检查 JAR 文件
    if [[ ! -f "$JAR_FILE" ]]; then
        log_error "JAR 文件不存在: $JAR_FILE"
        log_info "请先构建项目: mvn clean package"
        exit 1
    fi
    
    # 检查 Java 版本
    java_version=$(java -version 2>&1 | awk -F '"' '/version/ {print $2}')
    log_info "Java 版本: $java_version"
    
    # 检查 Hadoop 环境
    if command -v hadoop &> /dev/null; then
        hadoop_version=$(hadoop version 2>&1 | head -n1)
        log_info "Hadoop 版本: $hadoop_version"
    else
        log_warn "未检测到 Hadoop 环境"
    fi
    
    log_info "环境检查通过"
}

# 验证 HBase 连接
validate_hbase_connection() {
    local zk_quorum="$1"
    local table_name="$2"
    
    if [[ -z "$zk_quorum" || -z "$table_name" ]]; then
        log_warn "未指定 HBase 连接参数，跳过连接验证"
        return 0
    fi
    
    log_info "验证 HBase 连接: $zk_quorum"
    
    # 这里可以添加 HBase 连接验证逻辑
    # 例如使用 hbase shell 或者 Java 程序验证
    
    log_info "HBase 连接验证通过"
}

# 验证 HDFS 连接
validate_hdfs_connection() {
    local hdfs_path="$1"
    
    if [[ -z "$hdfs_path" ]]; then
        log_warn "未指定 HDFS 路径，跳过 HDFS 验证"
        return 0
    fi
    
    log_info "验证 HDFS 连接: $hdfs_path"
    
    # 尝试访问 HDFS
    if command -v hdfs &> /dev/null; then
        if hdfs dfs -test -d "$(dirname "$hdfs_path")" 2>/dev/null; then
            log_info "HDFS 连接验证通过"
        else
            log_warn "HDFS 路径不存在或无法访问: $(dirname "$hdfs_path")"
        fi
    else
        log_warn "未安装 HDFS 客户端，跳过验证"
    fi
}

# 验证 Kerberos 配置
validate_kerberos() {
    local principal="$1"
    local keytab_path="$2"
    
    if [[ -z "$principal" || -z "$keytab_path" ]]; then
        return 0
    fi
    
    log_info "验证 Kerberos 配置..."
    
    # 检查 keytab 文件
    if [[ ! -f "$keytab_path" ]]; then
        log_error "Keytab 文件不存在: $keytab_path"
        exit 1
    fi
    
    # 检查 kinit 命令
    if command -v kinit &> /dev/null; then
        if kinit -kt "$keytab_path" "$principal" 2>/dev/null; then
            log_info "Kerberos 认证验证通过"
            klist | head -3
        else
            log_error "Kerberos 认证失败"
            exit 1
        fi
    else
        log_warn "未安装 Kerberos 客户端工具"
    fi
}

# 构建 Flink 提交命令
build_flink_command() {
    local cmd="$FLINK_HOME/bin/flink run"
    
    # 添加 JAR 文件
    cmd="$cmd $JAR_FILE"
    
    # 添加主要参数
    if [[ -n "$ZK_QUORUM" ]]; then
        cmd="$cmd $ZK_QUORUM"
    fi
    
    if [[ -n "$TABLE_NAME" ]]; then
        cmd="$cmd $TABLE_NAME"
    fi
    
    if [[ -n "$HDFS_OUTPUT" ]]; then
        cmd="$cmd $HDFS_OUTPUT"
    fi
    
    if [[ -n "$CHECKPOINT_PATH" ]]; then
        cmd="$cmd $CHECKPOINT_PATH"
    fi
    
    if [[ -n "$PRINCIPAL" ]]; then
        cmd="$cmd $PRINCIPAL"
    fi
    
    if [[ -n "$KEYTAB_PATH" ]]; then
        cmd="$cmd $KEYTAB_PATH"
    fi
    
    echo "$cmd"
}

# 显示作业配置信息
show_job_config() {
    cat << EOF

${GREEN}=== 作业配置信息 ===${NC}
${BLUE}HBase 配置:${NC}
  Zookeeper Quorum: ${ZK_QUORUM:-未设置}
  HBase 表名: ${TABLE_NAME:-未设置}
  列族名: ${COLUMN_FAMILY:-info}

${BLUE}HDFS 配置:${NC}
  输出路径: ${HDFS_OUTPUT:-未设置}
  Checkpoint 路径: ${CHECKPOINT_PATH:-未设置}

${BLUE}Kerberos 配置:${NC}
  启用状态: ${ENABLE_KERBEROS:-false}
  主体: ${PRINCIPAL:-未设置}
  Keytab 路径: ${KEYTAB_PATH:-未设置}

${BLUE}并行度配置:${NC}
  全局并行度: ${PARALLELISM:-4}
  数据源并行度: ${SOURCE_PARALLELISM:-2}
  数据写入并行度: ${SINK_PARALLELISM:-2}

${BLUE}作业信息:${NC}
  JAR 文件: $JAR_FILE
  主类: $MAIN_CLASS
  Flink 版本: $(cat "$FLINK_HOME/conf/flink-conf.yaml" | grep -E "^#.*version.*:" | head -1 || echo "未知")

EOF
}

# 主函数
main() {
    # 初始化变量
    FLINK_HOME="${DEFAULT_FLINK_HOME}"
    CONFIG_FILE="${DEFAULT_CONFIG_FILE}"
    DRY_RUN=false
    VERBOSE=false
    VALIDATE_ONLY=false
    ENABLE_KERBEROS=false
    COLUMN_FAMILY="info"
    PARALLELISM=4
    SOURCE_PARALLELISM=2
    SINK_PARALLELISM=2
    
    # 解析命令行参数
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                print_usage
                exit 0
                ;;
            -c|--config)
                CONFIG_FILE="$2"
                shift 2
                ;;
            -f|--flink-home)
                FLINK_HOME="$2"
                shift 2
                ;;
            -j|--jar)
                JAR_FILE="$2"
                shift 2
                ;;
            --zk-quorum)
                ZK_QUORUM="$2"
                shift 2
                ;;
            --table-name)
                TABLE_NAME="$2"
                shift 2
                ;;
            --column-family)
                COLUMN_FAMILY="$2"
                shift 2
                ;;
            --enable-kerberos)
                ENABLE_KERBEROS=true
                shift
                ;;
            --principal)
                PRINCIPAL="$2"
                shift 2
                ;;
            --keytab-path)
                KEYTAB_PATH="$2"
                shift 2
                ;;
            --krb5-conf-path)
                KRB5_CONF_PATH="$2"
                shift 2
                ;;
            --hdfs-output)
                HDFS_OUTPUT="$2"
                shift 2
                ;;
            --checkpoint-path)
                CHECKPOINT_PATH="$2"
                shift 2
                ;;
            --parallelism)
                PARALLELISM="$2"
                shift 2
                ;;
            --source-parallelism)
                SOURCE_PARALLELISM="$2"
                shift 2
                ;;
            --sink-parallelism)
                SINK_PARALLELISM="$2"
                shift 2
                ;;
            --dry-run)
                DRY_RUN=true
                shift
                ;;
            --verbose)
                VERBOSE=true
                shift
                ;;
            --validate)
                VALIDATE_ONLY=true
                shift
                ;;
            *)
                log_error "未知参数: $1"
                print_usage
                exit 1
                ;;
        esac
    done
    
    # 检查环境
    check_environment
    
    # 如果只是验证模式
    if [[ "$VALIDATE_ONLY" == true ]]; then
        validate_hbase_connection "$ZK_QUORUM" "$TABLE_NAME"
        validate_hdfs_connection "$HDFS_OUTPUT"
        validate_kerberos "$PRINCIPAL" "$KEYTAB_PATH"
        log_info "环境验证完成"
        exit 0
    fi
    
    # 验证必需参数
    if [[ -z "$ZK_QUORUM" || -z "$TABLE_NAME" || -z "$HDFS_OUTPUT" ]]; then
        log_error "缺少必需参数: zk-quorum, table-name, hdfs-output"
        print_usage
        exit 1
    fi
    
    # 验证连接
    validate_hbase_connection "$ZK_QUORUM" "$TABLE_NAME"
    validate_hdfs_connection "$HDFS_OUTPUT"
    
    # 如果启用了 Kerberos
    if [[ "$ENABLE_KERBEROS" == true ]]; then
        validate_kerberos "$PRINCIPAL" "$KEYTAB_PATH"
    fi
    
    # 显示配置信息
    if [[ "$VERBOSE" == true ]]; then
        show_job_config
    fi
    
    # 构建命令
    FLINK_CMD=$(build_flink_command)
    
    # 显示信息
    log_info "准备提交 HBase 到 HDFS Parquet 作业"
    
    if [[ "$VERBOSE" == true ]]; then
        log_debug "完整命令: $FLINK_CMD"
    fi
    
    # 执行或显示命令
    if [[ "$DRY_RUN" == true ]]; then
        log_info "DRY RUN - 以下命令将被执行:"
        echo "$FLINK_CMD"
    else
        log_info "提交作业..."
        eval "$FLINK_CMD"
        
        if [[ $? -eq 0 ]]; then
            log_info "作业提交成功"
            log_info "可以使用 Flink Web UI 监控作业状态: http://localhost:8081"
            log_info "HDFS 输出路径: $HDFS_OUTPUT"
        else
            log_error "作业提交失败"
            exit 1
        fi
    fi
}

# 执行主函数
main "$@" 